package graph;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

public final class NodeHeap<T> {
    /**
     * 堆数组
     */
    private List<Node<T>> heap;

    /**
     * 存储源点到任意点的最短路径
     */
    private Map<Node<T>, Integer> weightMap;

    /**
     * 元素在堆中的索引
     * 为了支持修改已经堆中元素的权重，并将其正确堆化，需要记录元素在堆中的索引
     */
    private Map<Node<T>, Integer> indexMap;

    /**
     * indexMap中value为FREEZE_FLAG，表示当前节点的最短距离已被确定
     */
    private static final int FREEZE_FLAG = -1;

    /**
     * 表示当前节点不在堆中
     */
    private static final int NOT_EXISTS_FLAG = -2;

    /**
     * 堆弹出时的返回结果
     */
    public static final class NodeWithWeight<T> {
        private Node<T> node;
        private int weight;

        public NodeWithWeight(Node<T> node, int weight) {
            this.node = node;
            this.weight = weight;
        }

        public Node<T> getNode() {
            return node;
        }

        public void setNode(Node<T> node) {
            this.node = node;
        }

        public int getWeight() {
            return weight;
        }

        public void setWeight(int weight) {
            this.weight = weight;
        }
    }

    public NodeHeap() {
        this.heap = new ArrayList<>();
        this.weightMap = new HashMap<>();
        this.indexMap = new HashMap<>();
    }

    public boolean isEmpty() {
        return heap.isEmpty();
    }

    public int size() {
        return heap.size();
    }

    /**
     * 将节点加入到堆中
     *
     * @param node   节点
     * @param weight 权重
     * @return true表示当前操作更新了距离表，false反之
     */
    public boolean push(Node<T> node, Integer weight) {
        int index = indexMap.getOrDefault(node, NOT_EXISTS_FLAG);
        if (index == FREEZE_FLAG) {
            // 当前节点已经确定最短距离
            return false;
        }

        if (index == NOT_EXISTS_FLAG) {
            heap.add(node);
            index = heap.size() - 1;
            indexMap.put(node, index);
        }

        // 尝试更新距离表
        weightMap.put(node, Math.min(weightMap.getOrDefault(node, Integer.MAX_VALUE), weight));
        // 因为新权重一定小于等于之前权重，且是小根堆，因此这里只考虑向上堆化
        heapInsert(index);

        return true;
    }

    public NodeWithWeight<T> pop() {
        if (heap.isEmpty()) {
            throw new IndexOutOfBoundsException("heap is empty");
        }

        Node<T> top = heap.get(0);
        final int lastIndex = heap.size() - 1;
        swap(0, lastIndex);
        heap.remove(lastIndex);
        // 标记top节点已确定最短距离
        indexMap.put(top, FREEZE_FLAG);

        heapify(0);
        return new NodeWithWeight<>(top, weightMap.get(top));
    }

    /**
     * 将index位置的节点，向上堆化
     */
    private void heapInsert(int index) {
        while (weightMap.get(heap.get(index)) < weightMap.get(heap.get(parent(index)))) {
            swap(index, parent(index));
            index = parent(index);
        }
    }

    /**
     * 将index位置的节点向下堆化
     */
    private void heapify(int index) {
        int left = left(index);
        while (left < heap.size()) {
            int smaller = left + 1 < heap.size() ?
                    weightMap.get(heap.get(left)) < weightMap.get(heap.get(left + 1)) ? left : left + 1
                    : left;
            smaller = weightMap.get(heap.get(index)) < weightMap.get(heap.get(smaller)) ? index : smaller;
            if (smaller == index) break;
            swap(index, smaller);
            index = smaller;
            left = left(index);
        }
    }

    /**
     * 交换堆中i，j索引位置节点
     * 同步更新索引表
     */
    private void swap(int i, int j) {
        Node<T> iNode = heap.get(i);
        Node<T> jNode = heap.get(j);
        heap.set(i, jNode);
        heap.set(j, iNode);
        indexMap.put(iNode, j);
        indexMap.put(jNode, i);
    }

    private static int left(int index) {
        return index * 2 + 1;
    }

    private static int right(int index) {
        return index * 2 + 2;
    }

    private static int parent(int index) {
        return (index - 1) / 2;
    }
}
